#include "CNN/bgv_mnist_helper.hpp"
#include "intel_itt_wrapper.hpp"
#include "math.cpp"

INTEL_ITT_DOMAIN_CREATE(mnist_forward_test, "test.cnn.mnist.forward");
INTEL_ITT_STRING_HANDLE_CREATE(hcnn_conv, "HConv");
INTEL_ITT_STRING_HANDLE_CREATE(hcnn_sq1, "HSquare1");
INTEL_ITT_STRING_HANDLE_CREATE(hcnn_fc1, "HFC1");
INTEL_ITT_STRING_HANDLE_CREATE(hcnn_sq2_fc2, "HSquare2-HFC2");

static inline double log2_noise_bound(const helib::Ctxt &ctxt) {
  const double ln2 = NTL::log(NTL::xdouble(2l));
  return NTL::log(ctxt.totalNoiseBound()) / ln2;
}

static inline void trans_neg(NTL::mat_ZZ &a, long t) {
  for (int i = 0; i < a.NumRows(); i++)
    for (int j = 0; j < a.NumCols(); j++) {
      if (NTL::conv<int>(a[i][j]) > (t / 2)) {
        a[i][j] -= t;
      }
    }
}

static inline void trans_neg_single(NTL::ZZ &a, const NTL::ZZ &t) {
  if (a >= (t / 2))
    a -= t;
}

static inline void Tensor2Matrix(NTL::mat_ZZ &mat, int row, int col,
                                 const torch::Tensor &a, const NTL::ZZ &T,
                                 const NTL::ZZ &t) {
  mat.SetDims(row, col);
  for (int i = 0; i < row; i++)
    for (int j = 0; j < col; j++) {
      if (a[i][j].item<float>() >= 0)
        mat[i][j] = (NTL::conv<NTL::ZZ>(a[i][j].item<float>() * SCALE)) % t;
      else
        mat[i][j] =
            ((NTL::conv<NTL::ZZ>(a[i][j].item<float>() * SCALE)) + T) % t;
    }
}

static inline void Tensor2Matrix_conv(NTL::mat_ZZ &mat, int row, int col,
                                      const torch::Tensor &a, const NTL::ZZ &T,
                                      const NTL::ZZ &t) {
  mat.SetDims(row, col);
  for (int i = 0; i < row; i++)
    for (int j = 0; j < col; j++) {
      if (a[0][0].item<float>() >= 0)
        mat[i][j] = (NTL::conv<NTL::ZZ>(a[0][0].item<float>() * SCALE)) % t;
      else
        mat[i][j] =
            ((NTL::conv<NTL::ZZ>(a[0][0].item<float>() * SCALE)) + T) % t;
    }
  // std::cout << mat[0][0] << " " << a[0][0].item() << std::endl;
}

// clang-format off
static inline torch::Tensor Matrix2Tensor(const NTL::mat_ZZ &mat, int row, int col) {
  // clang-format on
  torch::Tensor a = torch::zeros({row, col});
  for (int i = 0; i < row; i++)
    for (int j = 0; j < col; j++)
      a[i][j] = NTL::conv<int>(mat[i][j]);
  return std::move(a);
}

helib::Ctxt bgv_hmnist_client::init_hx(const torch::Tensor &a, int asz0,
                                       int t_index) const {
  // input data
  NTL::mat_ZZ ptxt;
  Tensor2Matrix(ptxt, a.size(2), a.size(3), a[asz0][0], T_, ti_[t_index]);

  // Encryption
  auto ctxt = hmmccs_[t_index]->encrypt(ptxt, key);
  return std::move(ctxt);
}

helib::Ctxt bgv_hmnist_client::init_hconv(const torch::Tensor &a, int asz0,
                                          int asz1, int t_index) const {
  // input data
  NTL::mat_ZZ ptxt;
  Tensor2Matrix_conv(ptxt, 64, 64, a[asz0][asz1], T_, ti_[t_index]);
  // Encryption
  auto ctxt = hmmccs_[t_index]->encrypt(ptxt, key);
  // auto AB = hmmccs_[t_index]->decrypt(ctxt);
  // trans_neg(AB, NTL::conv<long>(t));
  // std::cout << a[asz0][asz1][0][0].item() << ", "
  //           << NTL::conv<float>(AB[0][0]) / 1000 << std::endl;
  return std::move(ctxt);
}

helib::Ctxt bgv_hmnist_client::init_hfc1(const torch::Tensor &a, int asz0,
                                         int t_index) const {
  // input data
  NTL::mat_ZZ ptxt;
  // std::cout << a.size(1) << " " << a.size(2) << std::endl;
  Tensor2Matrix(ptxt, a.size(1), a.size(2), a[asz0], T_, ti_[t_index]);

  // Encryption
  auto ctxt = hmmccs_[t_index]->encrypt(ptxt, key);
  return std::move(ctxt);
}

helib::Ctxt bgv_hmnist_client::init_hfc2(const torch::Tensor &a,
                                         int t_index) const {
  // input data
  NTL::mat_ZZ ptxt;
  Tensor2Matrix(ptxt, a.size(0), a.size(1), a, T_, ti_[t_index]);

  // Encryption
  auto ctxt = hmmccs_[t_index]->encrypt(ptxt, key);
  return std::move(ctxt);
}

bgv_hmnist_client::bgv_hmnist_client(const params<helib::BGV> &params,
                                     int crt_num)
    : crt_num_(crt_num) {
  hmmccs_.resize(crt_num);
  ti_.resize(crt_num);

  auto start_t = omp_get_wtime();
#pragma omp parallel for
  for (int i = 0; i < crt_num; i++) {
    hmmccs_[i] = new hypercube_hmmcc(params, MNIST_HYPERCUBE, true);
    ti_[i] = hmmccs_[i]->context()->getP();
  }

  T_ = NTL::to_ZZ(1l);
  contexts_.resize(crt_num);
  for (int i = 0; i < crt_num; i++) {
    contexts_[i] = hmmccs_[i]->context().get();
    hmmccs_[i]->report_context();
    T_ *= ti_[i];
  }
  auto end_t = omp_get_wtime();
  std::cout << "* Setup Helib Context           : " << (end_t - start_t) << " s"
            << std::endl;
}

void bgv_hmnist_client::init_hcnn(encrypted_model<helib::BGV> &ctxt_model,
                                  const torch::nn::Conv2d &conv1,
                                  const torch::nn::Linear &fc1,
                                  const torch::nn::Linear &fc2) const {
  auto init_hcnn_start_t = omp_get_wtime();
  auto b = init_conv_kernel(conv1->weight.data()); // [49,4,1,1]
  auto c = init_fc1_kernel(fc1->weight.data(), 4); // [4,64,64]
  auto d = init_fc2_kernel(fc2->weight.data());    // [10,64]

// init Convolution layer
#pragma omp parallel
#pragma omp single
  {
    auto start_t = omp_get_wtime();
    ctxt_model.conv.resize(crt_num_);
    for (int k = 0; k < crt_num_; k++) {
      ctxt_model.conv[k].resize(b.size(0));
      for (int i = 0; i < b.size(0); i++) {
        for (int j = 0; j < b.size(1); j++) {
          ctxt_model.conv[k][i].emplace_back(hmmccs_[k]->public_key());
        }
      }
    }
#pragma omp taskloop collapse(3) nogroup
    for (int k = 0; k < crt_num_; k++) {
      for (int i = 0; i < b.size(0); i++) {
        for (int j = 0; j < b.size(1); j++) {
          ctxt_model.conv[k][i][j] = std::move(init_hconv(b, i, j, k));
        }
      }
    }
    // for (int k = 0; k < crt_num_; k++) {
    //   for (int i = 0; i < b.size(0); i++) {
    //     for (int j = 0; j < b.size(1); j++) {
    //       auto AB = hmmccs_[k]->decrypt(ctxt_model.conv[k][i][j]);
    //       auto t = hmmccs_[k]->context()->getP();
    //       trans_neg_single(AB[0][0], NTL::conv<NTL::ZZ>(t));
    //       std::cout << b[i][j].item<float>() << ", "
    //                 << NTL::conv<NTL::xdouble>(AB[0][0]) / SCALE <<
    //                 std::endl;
    //     }
    //   }
    // }
    auto end_t = omp_get_wtime();
    std::cout << "  * Init convolution layer      : " << (end_t - start_t)
              << " s\n";

    // init fully connected layer1
    start_t = omp_get_wtime();
    ctxt_model.fc1.resize(crt_num_);
    for (int k = 0; k < crt_num_; k++) {
      for (int i = 0; i < c.size(0); i++) {
        ctxt_model.fc1[k].emplace_back(hmmccs_[k]->public_key());
      }
    }
#pragma omp taskloop collapse(2) nogroup
    for (int k = 0; k < crt_num_; k++) {
      for (int i = 0; i < c.size(0); i++) {
        ctxt_model.fc1[k][i] = std::move(init_hfc1(c, i, k));
      }
    }
    end_t = omp_get_wtime();
    std::cout << "  * Init fully connected layer1 : " << (end_t - start_t)
              << " s\n";

    // init fully connected layer2
    start_t = omp_get_wtime();

    for (int k = 0; k < crt_num_; k++) {
      ctxt_model.fc2.emplace_back(hmmccs_[k]->public_key());
    }
#pragma omp taskloop nogroup
    for (int k = 0; k < crt_num_; k++) {
      ctxt_model.fc2[k] = std::move(init_hfc2(d, k));
    }
    end_t = omp_get_wtime();
    std::cout << "  * Init fully connected layer2 : " << (end_t - start_t)
              << " s\n";
    auto init_hcnn_end_t = omp_get_wtime();
    std::cout << "* Init HCNN                     : "
              << (init_hcnn_end_t - init_hcnn_start_t) << " s\n";
  }
}

void bgv_hmnist_client::init_input(encrypted_input<helib::BGV> &ctxt_input,
                                   const torch::Tensor &filter,
                                   const torch::Tensor &x) const {
  auto a = im2matrix(x, filter, 3);
  std::vector<helib::Ctxt> tmp;
  std::vector<std::vector<helib::Ctxt>> tmp2;

#pragma omp parallel
#pragma omp single
  {
    // init input
    auto start_t = omp_get_wtime();
    ctxt_input.data.resize(crt_num_);
    for (int k = 0; k < crt_num_; k++) {
      ctxt_input.data[k].resize(a.size(0));
      for (int i = 0; i < a.size(0); i++) {
        for (int j = 0; j < a.size(1); j++) {
          ctxt_input.data[k][i].emplace_back(hmmccs_[k]->public_key());
        }
      }
    }
#pragma omp taskloop collapse(3)
    for (int k = 0; k < crt_num_; k++) {
      for (int i = 0; i < a.size(0); i++) {
        for (int j = 0; j < a.size(1); j++) {
          ctxt_input.data[k][i][j] = std::move(init_hx(a, i, k));
        }
      }
    }
    auto end_t = omp_get_wtime();
    std::cout << "* Init input                    : " << (end_t - start_t)
              << " s\n";
  }
}

void bgv_hmnist_client::recover_result(
    at::Tensor &out, const encrypted_output<helib::BGV> &in) const {
  auto start_t = omp_get_wtime();
  // CRT-Pre
  CH_Remainder crt[out.size(0)][out.size(1)][crt_num_];
  for (int k = 0; k < crt_num_; k++) {
    NTL::mat_ZZ res;
    hmmccs_[k]->decrypt(res, in.data[k]);
    for (int i = 0; i < out.size(0); i++)
      for (int j = 0; j < out.size(1); j++) {
        crt[i][j][k].result = NTL::conv<NTL::ZZ>(res[i][j]);
        crt[i][j][k].mod_num = NTL::conv<NTL::ZZ>(ti_[k]);
        // x.print();
      }
  }

  // CRT && Scale
  NTL::ZZ pow, base(SCALE);
  NTL::power(pow, base, 11);
  for (int i = 0; i < out.size(0); i++) {
    for (int j = 0; j < out.size(1); j++) {
      NTL::ZZ temp = Ch_remainder_theorem(crt[i][j], crt_num_);
      trans_neg_single(temp, T_);
      out.index({i, j}) = NTL::conv<float>(
          (NTL::conv<NTL::xdouble>(temp) / NTL::conv<NTL::xdouble>(pow)));

      // std::cout << temp << " "
      //           << NTL::conv<NTL::xdouble>(temp) /
      //                  NTL::conv<NTL::xdouble>(pow)
      //           << " " << out[i][j].item<float>() << " "
      //           << x[i][j].item<float>() << std::endl;
    }
  }
  out = out.transpose(0, 1);
  auto end_t = omp_get_wtime();
  std::cout << "CRT: " << (end_t - start_t) << " s" << std::endl;
}

bgv_hmnist_server::bgv_hmnist_server(
    const std::vector<helib::Context *> &contexts,
    const hmm_status<shmm_engine> &status, int crt_num)
    : crt_num_(crt_num) {
  metrics_.resize(crt_num);
  hmmes_.resize(crt_num);
  metrics_[0].hmm_encoding_time -= omp_get_wtime();
#pragma omp parallel for
  for (int i = 0; i < crt_num; i++) {
    hmmes_[i] = new hypercube_hmme(status);
    auto fc1 = hypercube_hmme::get_expected_mnk(MATRIX_DIM_M, MATRIX_DIM_N,
                                                MATRIX_DIM_K);
    hmmes_[i]->register_engine(contexts[i]->getEA(), fc1);
    auto fc2 =
        hypercube_hmme::get_expected_mnk(10l, MATRIX_DIM_N, MATRIX_DIM_K);
    hmmes_[i]->register_engine(contexts[i]->getEA(), fc2);
  }
  metrics_[0].hmm_encoding_time += omp_get_wtime();
}

void bgv_hmnist_server::forward_test(
    encrypted_output<helib::BGV> &out, const encrypted_input<helib::BGV> &in,
    const encrypted_model<helib::BGV> &ctxt_model) const {
  auto &metrics = const_cast<std::vector<metric> &>(metrics_);
  // Reset metrics
  for (auto &metric : metrics) {
    metric.reset();
  }

  auto ctxt = in.data;
  const auto dim0 = ctxt[0].size();
  const auto dim1 = ctxt[0][0].size();
  metrics[0].forword_test_time -= omp_get_wtime();
  omp_set_nested(1);
#pragma omp parallel
#pragma omp single
#pragma omp taskloop
  for (int k = 0; k < crt_num_; k++) {
    metrics[k].init_noise = log2_noise_bound(in.data[0][0][0]);
    // conv_hnmm
    metrics[k].conv_noise = metrics[k].init_noise;
    INTEL_ITT_TASK_BEGIN(mnist_forward_test, hcnn_conv); //! ITT - Begin
    metrics[k].hconv_time -= omp_get_wtime();
// #pragma omp parallel for collapse(2)
#pragma omp taskloop collapse(2)
    for (int i = 0; i < dim0; i++)
      for (int j = 0; j < dim1; j++) {
        ctxt[k][i][j] *= ctxt_model.conv[k][i][j];
      } //! omp taskloop: Implicit synchronization

      // reduction
// #pragma omp parallel for
#pragma omp taskloop
    for (int j = 0; j < dim1; j++) {
      for (int i = 1; i < dim0; i++) {
        ctxt[k][0][j] += ctxt[k][i][j];
      }
    } //! omp taskloop: Implicit synchronization

    // for (int k = 0; k < T_NUM; k++) {
    // for (int i = 0; i < dim0; i++) {
    //   for (int j = 0; j < dim1; j++) {
    //     auto AB = mnist_hmm[k].hmmcc_->decrypt(ctxt[k][i][j]);
    //     std::cout << AB << std::endl;
    //   }
    // }
    metrics[k].hconv_time += omp_get_wtime();
    INTEL_ITT_TASK_END(mnist_forward_test); //! ITT - End
    metrics[k].conv_noise =
        log2_noise_bound(ctxt[k][0][0]) - metrics[k].conv_noise;

    // sqr1
    metrics[k].sq1_noise = log2_noise_bound(ctxt[k][0][0]);
    INTEL_ITT_TASK_BEGIN(mnist_forward_test, hcnn_sq1); //! ITT - Begin
    metrics[k].hsq1_time -= omp_get_wtime();
// #pragma omp parallel for
#pragma omp taskloop
    for (int i = 0; i < dim1; i++) {
      ctxt[k][0][i].square();
    } //! omp taskloop: Implicit synchronization
    metrics[k].hsq1_time += omp_get_wtime();
    INTEL_ITT_TASK_END(mnist_forward_test); //! ITT - End
    metrics[k].sq1_noise =
        log2_noise_bound(ctxt[k][0][0]) - metrics[k].sq1_noise;

    // hfc1
    metrics[k].fc1_noise = log2_noise_bound(ctxt[k][0][0]);
    INTEL_ITT_TASK_BEGIN(mnist_forward_test, hcnn_fc1); //! ITT - Begin
    metrics[k].hfc1_time -= omp_get_wtime();
    // #pragma omp parallel for
    // #pragma omp taskloop
    for (int i = 0; i < dim1; i++) {
      hmmes_[k]->cgemm(MATRIX_TYPE::SQUARE, ctxt[k][0][i], ctxt_model.fc1[k][i],
                       ctxt[k][0][i]);
      if (i) {
        ctxt[k][0][0] += ctxt[k][0][i];
      }
    } //! omp taskloop: Implicit synchronization
    // for (int i = 1; i < dim1; i++) {
    //   ctxt[k][0][0] += ctxt[k][0][i];
    // }
    metrics[k].hfc1_time += omp_get_wtime();
    INTEL_ITT_TASK_END(mnist_forward_test); //! ITT - End
    metrics[k].fc1_noise =
        log2_noise_bound(ctxt[k][0][0]) - metrics[k].fc1_noise;

    // sqr2
    metrics[k].sq2_noise = log2_noise_bound(ctxt[k][0][0]);
    INTEL_ITT_TASK_BEGIN(mnist_forward_test, hcnn_sq2_fc2); //! ITT - Begin
    metrics[k].hsq2_time -= omp_get_wtime();
    ctxt[k][0][0].square();
    metrics[k].hsq2_time += omp_get_wtime();
    metrics[k].sq2_noise =
        log2_noise_bound(ctxt[k][0][0]) - metrics[k].sq2_noise;

    // hfc2
    metrics[k].fc2_noise = log2_noise_bound(ctxt[k][0][0]);
    metrics[k].hfc2_time -= omp_get_wtime();
    // hmmes_[k]->status()->set_clean_left(true);
    hmmes_[k]->cgemm(MATRIX_TYPE::MAX_K, ctxt[k][0][0], ctxt_model.fc2[k],
                     ctxt[k][0][0]);
    metrics[k].hfc2_time += omp_get_wtime();
    INTEL_ITT_TASK_END(mnist_forward_test); //! ITT - End
    metrics[k].fc2_noise =
        log2_noise_bound(ctxt[k][0][0]) - metrics[k].fc2_noise;
    metrics[k].final_noise = log2_noise_bound(ctxt[k][0][0]);
    metrics[k].left_cap = ctxt[k][0][0].capacity();
  }
  omp_set_nested(0);

  for (int k = 0; k < crt_num_; k++) {
    out.data.push_back(ctxt[k][0][0]);
  }

  metrics[0].forword_test_time += omp_get_wtime();
}

void bgv_hmnist_server::report_metrics(std::ostream &ss) const {
  // clang-format off
  ss << "================= Cipher CNN =================\n";
  ss << "* HMM Encoding       : " << metrics_[0].hmm_encoding_time << " s\n";
  ss << "* HCNN Time          : " << metrics_[0].forword_test_time << " s\n";
  for (auto i = 0; i < metrics_.size(); i++) {
  ss << "  * Round " << i << "          : \n";
  ss << "    * HConv time     : " << metrics_[i].hconv_time << " s\n";
  ss << "    * HSquare1 time  : " << metrics_[i].hsq1_time << " s\n";
  ss << "    * HFC1 time      : " << metrics_[i].hfc1_time << " s\n";
  ss << "    * HSquare2 time  : " << metrics_[i].hsq2_time << " s\n";
  ss << "    * HFC2 time      : " << metrics_[i].hfc2_time << " s\n";
  }
  ss << "* Final nosie bound  : " << metrics_[0].final_noise << std::endl;
  ss << "  * Init nosie bound : " << metrics_[0].init_noise << std::endl;
  ss << "  + HConv noise      : " << metrics_[0].conv_noise << std::endl;
  ss << "  + HSquare1 noise   : " << metrics_[0].sq1_noise << std::endl;
  ss << "  + HFC1 noise       : " << metrics_[0].fc1_noise << std::endl;
  ss << "  + HSquare2 noise   : " << metrics_[0].sq2_noise << std::endl;
  ss << "  + HFC2 noise       : " << metrics_[0].fc2_noise << std::endl;
  ss << "* Left capacity      : " << metrics_[0].left_cap << std::endl;
  ss << "==============================================\n";
  // clang-format on
}
